{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1c8e61d4-c237-4ca8-b7a7-7b9d2a0cdef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import jwt\n",
    "import pdfplumber\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import time\n",
    "import requests\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n",
    "from langchain.docstore.document import Document\n",
    "import torch\n",
    "from transformers import AutoModelForSequenceClassification, AutoTokenizer\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import re\n",
    "\n",
    "glm_key = os.getenv('glm_key')\n",
    "glm_key = '83e5bc58555d8bac289e27bac50f8afc.Khk1JjCxb8MJN8Mi'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dfe23e0-82d3-4b86-bc67-9aa98e7528cd",
   "metadata": {
    "tags": []
   },
   "source": [
    "# 读取数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "07a4ca37-4149-4116-8ca4-ee0e2709dac6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'question': '“前排座椅通风”的相关内容在第几页？', 'answer': '', 'reference': ''}, {'question': '\"关于车辆的儿童安全座椅固定装置，在哪一页可以找到相关内容？\"', 'answer': '', 'reference': ''}, {'question': '“打开前机舱盖”的相关信息在第几页？', 'answer': '', 'reference': ''}, {'question': '“打开前机舱盖”这个操作在哪一页？', 'answer': '', 'reference': ''}, {'question': '“查看行车记录仪视频”这一项内容在第几页？', 'answer': '', 'reference': ''}]\n",
      "354 页的说明书\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'欢迎\\n感谢您选择了具有优良安全性、舒适性、动力性和经济性的Lynk&Co领克汽车。\\n首次使用前请仔细、完整地阅读本手册内容，将有助于您更好地了解和使用车辆。\\n本手册中的所有资料均为出版时的最新资料，但本公司将对产品进行不断的改进和优化，您所购的车辆可能与本手册中的描述有所不同，请以实际\\n接收的车辆为准。\\n如您有任何问题，或需要预约服务，请拨打电话4006-010101联系我们。您也可以开车前往Lynk&Co领克中心。\\n在抵达之前，请您注意驾车安全。\\n©领克汽车销售有限公司'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "questions = json.load(open('data/questions.json', 'r', encoding='utf-8'))\n",
    "print(questions[:5])\n",
    "pdf = pdfplumber.open(\"data/初赛训练数据集.pdf\")\n",
    "print(len(pdf.pages),'页的说明书') # 页数\n",
    "pdf.pages[0].extract_text() # 读取第一页内容"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "04468462-61f6-4867-8893-4d40a37a6b72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>len</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>目录\\n前言 设置尾门开启角度..................................</td>\n",
       "      <td>3382</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>目录\\n组合仪表.........................................</td>\n",
       "      <td>4334</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>目录\\n电子驻车制动（EPB）..................................</td>\n",
       "      <td>3767</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>目录\\n系统设置.........................................</td>\n",
       "      <td>3432</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>前言\\n■ 保持清醒的驾驶状态，切勿在饮酒或服药后驾驶车辆，否则会影 体记录数据项的含义及用...</td>\n",
       "      <td>1109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>前言\\n原厂精装附件、选装装备和改装 在使用无线电设备时应当严格遵守以下规定：\\n■ 切勿自...</td>\n",
       "      <td>1030</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>180</th>\n",
       "      <td>驾驶辅助\\n高级智能驾驶 警告！\\n高级智能驾驶（LynkCo-Pilot）可以在0-130...</td>\n",
       "      <td>1024</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>186</th>\n",
       "      <td>驾驶辅助\\n■ 原车道与新车道之间存在较大的偏差道路。 警告！\\n■ 坑洼、突起、起伏路面。...</td>\n",
       "      <td>1210</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>219</th>\n",
       "      <td>泊车\\n泊车辅助摄像头局限性 ■ 如果泊车辅助雷达受到诸如尘土、雪或冰的阻塞，可能会导致泊\\...</td>\n",
       "      <td>1176</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>221</th>\n",
       "      <td>泊车\\n泊车辅助传感器位置 ■ 强烈的阳光、反光、光线较暗的场景下可能会使您难以看到视觉\\n...</td>\n",
       "      <td>1010</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>228</th>\n",
       "      <td>泊车\\n警告！ ■ 某些情况下泊车辅助传感器的视野会受到限制，系统检测到车\\n辆、行人或骑行...</td>\n",
       "      <td>1293</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                  text   len\n",
       "2    目录\\n前言 设置尾门开启角度..................................  3382\n",
       "3    目录\\n组合仪表.........................................  4334\n",
       "4    目录\\n电子驻车制动（EPB）..................................  3767\n",
       "5    目录\\n系统设置.........................................  3432\n",
       "11   前言\\n■ 保持清醒的驾驶状态，切勿在饮酒或服药后驾驶车辆，否则会影 体记录数据项的含义及用...  1109\n",
       "12   前言\\n原厂精装附件、选装装备和改装 在使用无线电设备时应当严格遵守以下规定：\\n■ 切勿自...  1030\n",
       "180  驾驶辅助\\n高级智能驾驶 警告！\\n高级智能驾驶（LynkCo-Pilot）可以在0-130...  1024\n",
       "186  驾驶辅助\\n■ 原车道与新车道之间存在较大的偏差道路。 警告！\\n■ 坑洼、突起、起伏路面。...  1210\n",
       "219  泊车\\n泊车辅助摄像头局限性 ■ 如果泊车辅助雷达受到诸如尘土、雪或冰的阻塞，可能会导致泊\\...  1176\n",
       "221  泊车\\n泊车辅助传感器位置 ■ 强烈的阳光、反光、光线较暗的场景下可能会使您难以看到视觉\\n...  1010\n",
       "228  泊车\\n警告！ ■ 某些情况下泊车辅助传感器的视野会受到限制，系统检测到车\\n辆、行人或骑行...  1293"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = [{'text': page.extract_text(), 'len': len(page.extract_text())} for page in pdf.pages]\n",
    "\n",
    "df = pd.DataFrame(data)\n",
    "df[df['len'] > 1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "75f7d6d5-3952-4ee1-b7f4-57e6b3a13afb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'泊车\\n警告！ - 某些情况下泊车辅助传感器的视野会受到限制，系统检测到车\\n辆、行人或骑行者的时间会比预期延后或完全检测不到车辆、行\\n- PEB仅是一种辅助功能，无法在所有情况下探测到所有行人、骑 人或骑行者。\\n行者或车辆。驾驶员需始终对正确驾驶负责，并需要保持安全距 - 当泊车辅助传感器被遮挡或功能受限时低速紧急制动辅助系统性\\n离。 能可能会下降甚至不可用。\\n- PEB只能提供报警及制动辅助，驾驶员需要时刻保持警惕并始终 - 受泊车辅助传感器安装位置影响，当车辆以非直线行驶时，处于\\n对车辆的安全驾驶负有责任，请遵守现行法律和交通法规。 行进方向的车轮轨迹外的障碍物可能无法被探测到。\\n- 通常PEB在后台工作，不会被驾驶员察觉，当系统识别到倒车有 - 湿滑路面上，车辆制动距离加长，这会降低PEB的防碰撞性能。\\n危险时，会采取制动来保护乘员。由于系统性能限制，可能存在\\n误触发，驾驶员必须始终密切注意周围环境。\\n- 泊车辅助传感器无法识别到超出探测能力的障碍物（如停车场起 说明！\\n停杆、树、草、细柱子、铁丝网、铁链、绳索、墙柱、栅栏、低\\n矮障碍物、吸波物体、反光物体等），可能会发生误报、漏报障 □ 当车速大于9km/h或低于2km/h时PEB功能自动退出。\\n碍物的情况。驾驶员要时刻注意周围情况，避免发生事故或造成\\n□ PEB激活时，您可以通过踩下制动踏板或者加速踏板解除PEB。\\n物品损坏。\\n- 系统可能无法及时对横穿或迎面而来的移动物体（行人、动物、 □ 当PEB激活后的30秒内未进行车辆其他操作，系统将会自动开\\n车辆等）或车辆侧面的物体做出反应。 启电子驻车制动。\\n- 请注意泊车辅助传感器不可能在所有情况下都探测到后方的危险 □ PEB功能激活/退出后，仪表显示屏上会显示相关信息。\\n障碍物。恶劣的天气条件，如雨、雪、雾等，会导致系统性能下\\n降，在此种情况下部分目标将无法被系统探测或探测不及时。\\n- 某些场景会对泊车辅助传感器的探测造成影响，如有防护栏的道\\n路、隧道内、前方车辆驶入/驶出、急转弯道路等。 泊车紧急制动局限性\\n- 系统不会对动物、小型车辆（如三轮车）、外表不规则车辆、骑 在以下情况下，PEB将无法正常使用：\\n行者、迎面而来及横穿的车辆进行及时反应。\\n- 快速移动的行人。\\n- 任何辅助系统都无法保证在任何情况下均可100%正常运行。因\\n- 特殊类型的行人（半蹲、坐轮椅、骑自行车等）。\\n此，请不要以测试PEB性能好坏为目的，倒车开向人或物体。否\\n则可能引发事故，导致人员伤亡。 - 在光照弱的环境下。\\n- 在复杂的行驶状况下，PEB可能会进行不必要的制动。例如在建 - 在地面反光严重的地下停车场环境下。\\n筑工地、铁轨处、道路窨井盖处、地下车库、车辆前方存在喷洒 - 当遇到恶劣天气时（如：暴雨天气）。\\n或溅起的水花时。\\n- 当泊车辅助传感器或摄像头表面被脏污覆盖或者系统发生故障。\\n- 对于系统识别的有效目标，根据车辆、行人、骑行者、场景、路\\n况的不同，本系统并非总能达到相同的性能水平。\\n关于更多PEB局限性，请参见泊车辅助系统章节。\\n229'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.iloc[228]['text'].replace('■', '-')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8bc8bdfb-290e-4910-9dbb-175c8ab2f07d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# rerank_model = SentenceTransformer('D:/env/bert_model/BAAI/bge-reranker-base', device='cuda' if torch.cuda.is_available() else 'cpu')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a0ad879-0657-4459-b31a-baf037cc0980",
   "metadata": {},
   "source": [
    "# 构建向量库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aade85b8-5760-40f9-94bf-3cb760ca5394",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = HuggingFaceEmbeddings(model_name='D:/env/bert_model/BAAI/bge-large-zh-v1.5')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3e0098ff-e016-4db0-b6a4-bacc3fd62b2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████| 354/354 [00:00<00:00, 58987.87it/s]\n"
     ]
    }
   ],
   "source": [
    "def pages_2_docs(pdf_pages):\n",
    "    docs = []\n",
    "    for i in tqdm(range(len(pdf_pages))):\n",
    "        context = pdf_pages[i].extract_text().replace('■', '-')\n",
    "        if len(context) < 1500:\n",
    "            if len(context) > 750:\n",
    "                split_index = int(len(context)*0.5)\n",
    "                context_list = [context[:split_index], context[split_index:]]\n",
    "            else:\n",
    "                context_list = [context]\n",
    "            for c in context_list:\n",
    "                docs.append(Document(page_content='【说明书的第{}页】'.format(i) + c, metadata={\"page_count\": 'page_{}'.format(i)}))\n",
    "    return docs\n",
    "\n",
    "# 生成向量库文档\n",
    "docs =pages_2_docs(pdf.pages)\n",
    "# 构建向量库\n",
    "vector_store = FAISS.from_documents(docs, embeddings)\n",
    "# 本地化存储\n",
    "vector_store.save_local('data/vector')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd693ad3-2194-4e11-9435-292b8ec02329",
   "metadata": {},
   "source": [
    "# 预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "138a751b-cbf8-4cd8-be37-cb8d2b54a781",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# chatglm\n",
    "def generate_token(apikey: str, exp_seconds: int):\n",
    "    try:\n",
    "        id, secret = apikey.split(\".\")\n",
    "    except Exception as e:\n",
    "        raise Exception(\"invalid apikey\", e)\n",
    "\n",
    "    payload = {\n",
    "        \"api_key\": id,\n",
    "        \"exp\": int(round(time.time() * 1000)) + exp_seconds * 1000,\n",
    "        \"timestamp\": int(round(time.time() * 1000)),\n",
    "    }\n",
    "    return jwt.encode(\n",
    "        payload,\n",
    "        secret,\n",
    "        algorithm=\"HS256\",\n",
    "        headers={\"alg\": \"HS256\", \"sign_type\": \"SIGN\"},\n",
    "    )\n",
    "\n",
    "def get_llm(question, token):\n",
    "    url = \"https://open.bigmodel.cn/api/paas/v4/chat/completions\"\n",
    "    headers = {\n",
    "      'Content-Type': 'application/json',\n",
    "      'Authorization': token\n",
    "    }\n",
    "\n",
    "    data = {\n",
    "        \"model\": \"glm-3-turbo\",\n",
    "        \"messages\": [{\"role\": \"user\", \"content\": question}]\n",
    "    }\n",
    "\n",
    "    response = requests.post(url, headers=headers, json=data)\n",
    "    return response.json()\n",
    "\n",
    "\n",
    "# 初始化大模型\n",
    "token = generate_token(glm_key, 3600)\n",
    "# 加载相关模型\n",
    "embeddings =HuggingFaceEmbeddings(model_name='D:/env/bert_model/BAAI/bge-large-zh-v1.5')\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained('D:/env/bert_model/BAAI/bge-reranker-base')\n",
    "rerank_model = AutoModelForSequenceClassification.from_pretrained('D:/env/bert_model/BAAI/bge-reranker-base')\n",
    "rerank_model.cuda()\n",
    "rerank_model.eval()\n",
    "# 加载向量库\n",
    "vector_store = FAISS.load_local('data/vector', embeddings)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ba49ad58-e0ec-4683-841c-7f9e9045830c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_llm(question, key):\n",
    "    url = \"https://open.bigmodel.cn/api/paas/v4/chat/completions\"\n",
    "\n",
    "    headers = {\n",
    "      'Content-Type': 'application/json',\n",
    "      'Authorization': key\n",
    "    }\n",
    "    data = {\n",
    "  \"model\": \"glm-4\",\n",
    "  \"messages\": [{\n",
    "      \"role\": \"user\",\n",
    "      \"content\": question\n",
    "    }]}\n",
    "    response = requests.post(url, headers=headers, json=data)\n",
    "    return response.json()\n",
    "token = generate_token(glm_key, 3600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "23274600-89c3-4b85-818d-63c266c96e3a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 301/301 [32:16<00:00,  6.43s/it]\n"
     ]
    }
   ],
   "source": [
    "llm_url = 'http://192.168.4.134:8008/v1/chat/completions'\n",
    "prompt = \"\"\"你是一个汽车专家，帮我结合给定的资料，回答一个问题。如果问题无法从资料中获得，请输出结合给定的资料，无法回答问题。\n",
    "【资料】\n",
    "{}\n",
    "\n",
    "【问题】{}\n",
    "\"\"\"\n",
    "answer_list = []\n",
    "# 遍历问题\n",
    "for question in tqdm(questions):\n",
    "    # 向量库召回\n",
    "    docs = vector_store.similarity_search(question['question'],k=5)\n",
    "    \n",
    "    # 提取文档列表\n",
    "    content_list = [doc.page_content for doc in docs]\n",
    "    # 数据重排\n",
    "    pairs = []\n",
    "    for doc in docs:\n",
    "        pairs.append([question['question'], re.sub(r'【[^】]*】', '', doc.page_content)][:150])\n",
    "    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)\n",
    "    with torch.no_grad():\n",
    "        inputs = {key: inputs[key].cuda() for key in inputs.keys()}\n",
    "        scores = rerank_model(**inputs, return_dict=True).logits.view(-1, ).float()\n",
    "    \n",
    "    index = scores.cpu().numpy().argmax()\n",
    "    question['reference'] = docs[index].metadata['page_count']\n",
    "    \n",
    "    r = docs[index].page_content\n",
    "\n",
    "    resp = get_llm(prompt.format(r, question['question']), token)['choices'][0]['message']['content']\n",
    "    answer_list.append({'question': question['question'], 'answer': resp, 'reference': docs[0].metadata['page_count']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "358d10e8-9f04-4901-9ea0-f4e204dddcbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://competition.coggle.club/\n",
    "with open('result/submitV5.json', 'w', encoding='utf8') as up:\n",
    "    json.dump(answer_list, up, ensure_ascii=False, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b650ca07-8da5-40af-9740-7533028c317a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "安全出行\n",
      "儿童座椅固定装置\n",
      "儿童安全\n",
      "儿童乘车时，请务必时刻确保儿童乘坐安全。\n",
      "婴儿和儿童必须乘坐在适合其年龄和体型的儿童安全座椅中。\n",
      "01儿童锁已停用，可以从车内和车外打开后门。\n",
      "02儿童锁已启用，不能从车内打开后门。\n",
      "儿童锁位于后门的后缘，后门打开时方可触及儿童锁。\n",
      "可单独启用后排儿童锁：\n",
      "使用可分开的机械钥匙，顺时针转动控制器90°（左后门）。 儿童座椅警告标贴\n",
      "不得在受正面安全气囊保护（激活状态下）的座位上使用后向儿童约\n",
      "束系统。\n",
      "说明！\n",
      "□ 对于右后门的儿童锁，逆时针转动钥匙90°，启用儿童锁。\n",
      "按相反方向转动相应控制器，可以停用儿童锁。\n",
      "122\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "\n",
    "# 原始字符串\n",
    "text = \"【说明书的第121页】安全出行\\n儿童座椅固定装置\\n儿童安全\\n儿童乘车时，请务必时刻确保儿童乘坐安全。\\n婴儿和儿童必须乘坐在适合其年龄和体型的儿童安全座椅中。\\n01儿童锁已停用，可以从车内和车外打开后门。\\n02儿童锁已启用，不能从车内打开后门。\\n儿童锁位于后门的后缘，后门打开时方可触及儿童锁。\\n可单独启用后排儿童锁：\\n使用可分开的机械钥匙，顺时针转动控制器90°（左后门）。 儿童座椅警告标贴\\n不得在受正面安全气囊保护（激活状态下）的座位上使用后向儿童约\\n束系统。\\n说明！\\n□ 对于右后门的儿童锁，逆时针转动钥匙90°，启用儿童锁。\\n按相反方向转动相应控制器，可以停用儿童锁。\\n122\"\n",
    "\n",
    "# 使用正则表达式替换【】内的内容为空\n",
    "cleaned_text = re.sub(r'【[^】]*】', '', text)\n",
    "\n",
    "print(cleaned_text)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0329a959-5b47-4bdc-92f9-38c1cbdba4b8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
